import unittest

from mysql_db_class import MySQLDB


class MySQLDBTest(unittest.TestCase):
    def setUp(self):
        self.db = MySQLDB('127.0.0.1', 'root', '123456', 'b23015114_zty')
        self.create_test_table()

    def tearDown(self):
        self.drop_test_table()
        self.db.disconnect()

    def create_test_table(self):
        with self.db as mysql_db:
            mysql_db.execute_update("""
                CREATE TABLE IF NOT EXISTS test_table (
                    id INT AUTO_INCREMENT PRIMARY KEY,
                    name VARCHAR(50) NOT NULL,
                    age INT
                )
            """)

    def drop_test_table(self):
        with self.db as mysql_db:
            mysql_db.execute_update("DROP TABLE IF EXISTS test_table")

    def test_connect(self):
        with self.db as mysql_db:
            self.assertIsNotNone(mysql_db.connection)
            self.assertIsInstance(mysql_db.connection, pymysql.connections.Connection)

    def test_execute_query(self):
        with self.db as mysql_db:
            # 插入测试数据
            mysql_db.execute_update("INSERT INTO test_table (name, age) VALUES (%s, %s)", ("Alice", 30))
            mysql_db.execute_update("INSERT INTO test_table (name, age) VALUES (%s, %s)", ("Bob", 25))

            # 执行查询并验证结果
            results = mysql_db.execute_query("SELECT * FROM test_table WHERE age > %s", (20,))
            self.assertEqual(len(results), 2)
            self.assertIn({"id": 1, "name": "Alice", "age": 30}, results)
            self.assertIn({"id": 2, "name": "Bob", "age": 25}, results)

    def test_execute_update(self):
        with self.db as mysql_db:
            # 插入测试数据
            mysql_db.execute_update("INSERT INTO test_table (name, age) VALUES (%s, %s)", ("Charlie", 40))

            # 更新数据并验证
            mysql_db.execute_update("UPDATE test_table SET age = %s WHERE name = %s", (45, "Charlie"))
            result = mysql_db.execute_query("SELECT * FROM test_table WHERE name = %s", ("Charlie",))
            self.assertEqual(result[0]['age'], 45)

if __name__ == '__main__':
    unittest.main()